{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43871e67",
   "metadata": {},
   "outputs": [],
   "source": [
    "%env CUDA_VISIBLE_DEVICES=6\n",
    "%env TRANSFORMERS_CACHE=/mnt/LLM/hub\n",
    "%env OMP_NUM_THREADS=16\n",
    "\n",
    "import os\n",
    "import sys\n",
    "sys.path.insert(0, '..')\n",
    "\n",
    "import time\n",
    "import random\n",
    "from tqdm.auto import trange\n",
    "import ipynbname  # pip install ipynbname\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import transformers\n",
    "\n",
    "from src.aq import QuantizedWeight\n",
    "\n",
    "\n",
    "torch.set_num_threads(min(16, saved_num_threads))\n",
    "torch.backends.cudnn.allow_tf32 = False\n",
    "torch.backends.cuda.matmul.allow_tf32 = False\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "input_loading_dir = '/DIR/layer10.self_attn.q_proj.input_activation.pt'\n",
    "num_codebooks = 2\n",
    "nbits_per_codebook = 12\n",
    "out_group_size = 1\n",
    "in_group_size = 8\n",
    "batch_size = 16384\n",
    "beam_size = 1\n",
    "beam_search_epochs = 100\n",
    "sparsity_regularizer = 0\n",
    "print_frequency = 10\n",
    "scale_nbits = 0    # 0 means no scales, 16 means no compression;\n",
    "codebook_values_nbits = 16  # less than 16 means we quantize codebooks as well\n",
    "init_max_iter = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7e6156a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "\n",
    "os.environ[\"WANDB_NOTEBOOK_NAME\"] = os.path.join(os.getcwd(), ipynbname.name() + \".ipynb\")\n",
    "\n",
    "# start a new wandb run to track this script\n",
    "run = wandb.init(\n",
    "    # set the wandb project where this run will be logged\n",
    "    dir=os.getcwd(),\n",
    "    project=\"\",\n",
    "    entity = \"\",\n",
    "    save_code=True,\n",
    "    name = f\"{ipynbname.name()}_AQ_{num_codebooks=}_{out_group_size=}_{in_group_size=}_{nbits_per_codebook=}_{beam_search_epochs=}\",\n",
    "    settings=wandb.Settings(code_dir=\".\"),\n",
    "    # track hyperparameters and run metadata\n",
    "    config={\n",
    "    \"num_codebooks\" : num_codebooks,\n",
    "    \"out_group_size\": out_group_size,\n",
    "    \"in_group_size\": in_group_size,\n",
    "    \"group_size\" : out_group_size * in_group_size,\n",
    "    \"batch_size\" : batch_size,\n",
    "    \"beam_size\" : beam_size,\n",
    "    \"nbits_per_codebook\" : nbits_per_codebook,\n",
    "    \"codebook_values_nbits\": codebook_values_nbits,\n",
    "    \"scale_nbits\": scale_nbits,\n",
    "    \"beam_search_epochs\": beam_search_epochs,\n",
    "    \"sparsity_regularizer\": sparsity_regularizer,\n",
    "    \"init_max_iter\": init_max_iter,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbf7e1f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = transformers.AutoModelForCausalLM.from_pretrained(\n",
    "    \"meta-llama/Llama-2-70b-hf\", torch_dtype='auto', low_cpu_mem_usage=True)\n",
    "\n",
    "X = torch.load(input_loading_dir, map_location='cpu').float().flatten(0, -2)\n",
    "reference_weight = model.model.layers[10].self_attn.q_proj.weight.detach().to(device).float()\n",
    "\n",
    "XTX = torch.zeros(X.shape[-1], X.shape[-1], device=device, dtype=torch.float64)\n",
    "for i in range(0, len(X), batch_size):\n",
    "    x_batch = X[i: i + batch_size].cuda().double()\n",
    "    XTX.addmm_(x_batch.T, x_batch, alpha=1/len(X))\n",
    "    del x_batch\n",
    "XTX = XTX.float()\n",
    "del X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e8656f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "quantized_weight = QuantizedWeight(\n",
    "    XTX=XTX, reference_weight=reference_weight, num_codebooks=num_codebooks,\n",
    "    nbits_per_codebook=nbits_per_codebook, scale_nbits=scale_nbits, \n",
    "    out_group_size=out_group_size, in_group_size=in_group_size,\n",
    "    verbose=True, max_iter=init_max_iter,   # faster init, not tested\n",
    ")\n",
    "run.log({\"Avg_bits\": quantized_weight.estimate_nbits_per_parameter()})\n",
    "print(\"AVG bits:\", quantized_weight.estimate_nbits_per_parameter())\n",
    "opt = torch.optim.Adam(quantized_weight.parameters(), lr=1e-4, betas=(0.0, 0.95), amsgrad=True)\n",
    "\n",
    "for epoch in range(1000):\n",
    "    start = time.perf_counter()\n",
    "    delta_weight = (quantized_weight() - reference_weight).double()\n",
    "    loss = (delta_weight @ XTX.double()).flatten() @ delta_weight.flatten() / len(delta_weight)\n",
    "    opt.zero_grad()\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "    \n",
    "    run.log({'loss':loss.item()}, step=epoch)\n",
    "    \n",
    "    if epoch % print_frequency == 0:\n",
    "        print(f\"loss={loss.item():.10f}\\t\",\n",
    "              f\"time_on_epoch {epoch} = {time.perf_counter() - start}\")\n",
    "    if (epoch + 1) % beam_search_epochs == 0:\n",
    "        quantized_weight.beam_search_update_codes_(\n",
    "            XTX, reference_weight, beam_size=beam_size, sparsity_regularizer=sparsity_regularizer,\n",
    "            dim_rng=random.Random(), verbose=True)\n",
    "\n",
    "        if sparsity_regularizer != 0:\n",
    "            sparsity_rate = ((quantized_weight.codes == 0).sum() / quantized_weight.codes.numel()).item()\n",
    "            print(f\"Sparsity rate {sparsity_rate:.5f}\")\n",
    "            run.log({'sparsity rate': sparsity_rate}, step=epoch)\n",
    "            mean_code_nbits = sum(get_mean_nbits_by_codebook(quantized_weight.codes)) / num_codebooks\n",
    "            print(f\"mean_code_nbits {mean_code_nbits:.5f}\")\n",
    "            run.log({'Mean codebook length nbits': mean_code_nbits}, step=epoch)\n",
    "            if in_group_size > 1 and out_group_size > 1:\n",
    "                curr_avg_bits  = calc_avg_bits(num_codebooks, 1, mean_code_nbits,\n",
    "                                     nbits_per_codebook, in_features, out_features, scale_nbits)\n",
    "                run.log({\"Avg_bits\": curr_avg_bits}, step=epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5702df5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
